import argparse
import torch.onnx
from models.refinedet import build_refinedet
from data import VOCAnnotationTransform, VOCDetection, BaseTransform

def pth2onnx(input_file, output_file):

    net = build_refinedet('test', 320, 21)
    net.load_state_dict(torch.load(input_file, map_location='cpu'))
    net.eval()
 
    input_names = ["image"]
    output_names = ["scores_boxes"]
    dynamic_axes = {'image':{0:'-1'}, 'scores_boxes':{0:'-1'}}

    dataset = VOCDetection(args.voc_root, [('2007', 'test')],
                           BaseTransform(int(320), (104, 117, 123)),
                           VOCAnnotationTransform())
    img, _, _, _ = dataset.pull_item(0)
    img = img.unsqueeze(0)

    torch.onnx.export(net, img, output_file, input_names=input_names, dynamic_axes=dynamic_axes,
                      output_names=output_names, opset_version=11, verbose=True)


if __name__ == '__main__':

    parser = argparse.ArgumentParser(
    description='Single Shot MultiBox Detector export onnx')
    parser.add_argument('--trained_model',
                        default='RefineDet320_VOC_final.pth', type=str,
                        help='Trained state_dict file path to open')
    parser.add_argument('--voc_root', default='/root/datasets/VOCdevkit/',
                        help='Location of VOC root directory')
    args = parser.parse_args()

    input_file = args.trained_model
    output_file = r'RefineDet320_VOC_final.onnx'
    pth2onnx(input_file, output_file)
